package com.l2qq.vben.auth;

import cn.dev33.satoken.dao.SaTokenDaoDefaultImpl;
import cn.dev33.satoken.session.SaSession;
import cn.dev33.satoken.session.SaTerminalInfo;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import com.jfinal.plugin.activerecord.Db;
import com.jfinal.plugin.activerecord.Record;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class PgSaTokenDao extends SaTokenDaoDefaultImpl {
    // 自定义会话存储（独立于父类存储）
    private final Map<String, SaSession> sessionMap = new ConcurrentHashMap<>();
    @Override
    public void init() {
        super.init();
        loadFromDatabase();
    }

    @Override
    public void destroy() {
        saveToDatabase();
        super.destroy();
    }

    private void loadFromDatabase() {
        // 1. 加载Token数据
        List<Record> tokenRecords = Db.find("SELECT * FROM sa_token_store");
        for (Record r : tokenRecords) {
            String key = r.getStr("key");
            String value = r.getStr("value");
            long expireTime = r.getLong("expire_time");
            long timeout = calculateTimeout(expireTime);
            this.set(key, value, timeout);
        }

        // 2. 加载Session数据
        List<Record> sessionRecords = Db.find("SELECT * FROM sa_token_session");
        for (Record r : sessionRecords) {
            String sessionId = r.getStr("id");
            String sessionJson = r.getStr("value");
            long expireTime = r.getLong("expire_time");
            long timeout = calculateTimeout(expireTime);

            SaSession session = JSON.parseObject(sessionJson, SaSession.class);
            if (session != null) {
                this.setSession(session, timeout);
            }
        }
    }

    private void saveToDatabase() {
        // 使用事务确保原子性
        Db.tx(() -> {
            // 1. 保存Token数据
            Db.update("TRUNCATE TABLE sa_token_store");
            List<Record> tokenRecords = new ArrayList<>();
            for (String key : timedCache.keySet()) {
                if (isSessionKey(key)) continue; // 跳过session数据

                String value = get(key);
                long expireTime = getTimeout(key);

                if (expireTime != NOT_VALUE_EXPIRE) {
                    tokenRecords.add(new Record()
                        .set("key", key)
                        .set("value", value)
                        .set("expire_time", toDatabaseExpire(expireTime))
                    );
                }
            }
            Db.batchSave("sa_token_store", tokenRecords, 500);

            // 2. 保存Session数据
            Db.update("TRUNCATE TABLE sa_token_session");
            List<Record> sessionRecords = new ArrayList<>();
            for (String key : timedCache.keySet()) {
                if (!isSessionKey(key)) continue;

                SaSession session = getSession(key);
                if (session != null) {
                    long expireTime = getTimeout(key);

                    sessionRecords.add(new Record()
                        .set("id", key)
                        .set("value", JSON.toJSONString(session))
                        .set("expire_time", toDatabaseExpire(expireTime))
                    );
                }
            }
            Db.batchSave("sa_token_session", sessionRecords, 500);

            return true;
        });
    }

    // 辅助方法保持不变...

    private long calculateTimeout(long expireTime) {
        if (expireTime == NEVER_EXPIRE) {
            return NEVER_EXPIRE;
        }
        long current = System.currentTimeMillis();
        return expireTime > current ? (expireTime - current) / 1000 : NOT_VALUE_EXPIRE;
    }

    private void loadSessionData() {
        List<Record> sessionRecords = Db.find("SELECT * FROM sa_token_session");
        for (Record r : sessionRecords) {
            SaSession session = deserializeSession(r);
            if (session != null) {
                sessionMap.put(session.getId(), session);

                // 设置会话超时
                long expireTime = r.getLong("expire_time");
                if (expireTime != NEVER_EXPIRE) {
                    this.updateSessionTimeout(session.getId(), calculateTimeout(expireTime));
                }
            }
        }
        System.out.println("session Loaded " + sessionRecords.size() + " sessions from database");
    }


    private SaSession deserializeSession(Record r) {
        try {
            SaSession session = new SaSession(r.getStr("id"));

            // 设置基本属性
            session.setType(r.getStr("type"));
            session.setLoginType(r.getStr("login_type"));
            session.setLoginId(r.getStr("login_id"));
            session.setToken(r.getStr("token"));

            // 设置历史终端计数
            Integer historyCount = r.getInt("history_terminal_count");
            if (historyCount != null) {
                session.setHistoryTerminalCount(historyCount);
            }

            // 设置终端列表
            String terminalListJson = r.getStr("terminal_list");
            if (terminalListJson != null && !terminalListJson.isEmpty()) {
                List<SaTerminalInfo> terminalList = JSON.parseArray(terminalListJson, SaTerminalInfo.class);
                session.setTerminalList(terminalList);
            }

            // 设置数据映射
            String dataMapJson = r.getStr("data_map");
            if (dataMapJson != null && !dataMapJson.isEmpty()) {
                Map<String, Object> dataMap = JSON.parseObject(dataMapJson,
                        new TypeReference<ConcurrentHashMap<String, Object>>() {});
                session.refreshDataMap(dataMap);
            }

            return session;
        } catch (Exception e) {
            System.err.println("Error deserializing session: " + e.getMessage());
            return null;
        }
    }

    /**
     * 判断是否为 session 键
     */
    private boolean isSessionKey(String key) {
        return key != null &&
                (key.startsWith("satoken:login:session:") ||
                        key.startsWith("satoken:login:token-session:"));
    }

    private long toDatabaseExpire(long timeoutSeconds) {
        if (timeoutSeconds == NEVER_EXPIRE) {
            return NEVER_EXPIRE;
        }
        return System.currentTimeMillis() + timeoutSeconds * 1000;
    }
}